# coding: utf-8
import os
import ahocorasick


class QuestionClassifier:
    def __init__(self):
        cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        # 　特征词路径
        self.director_path = os.path.join(cur_dir, 'dict/director.txt')
        self.actor_path = os.path.join(cur_dir, 'dict/actor.txt')
        self.less_path = os.path.join(cur_dir, 'dict/less.txt')
        self.more_path = os.path.join(cur_dir, 'dict/more.txt')
        self.equal_path = os.path.join(cur_dir, 'dict/equal.txt')
        self.movie_path = os.path.join(cur_dir, 'dict/movie.txt')
        self.type_path = os.path.join(cur_dir, 'dict/type.txt')
        self.date_path = os.path.join(cur_dir, 'dict/date.txt')
        self.star_path = os.path.join(cur_dir, 'dict/star.txt')
        self.time_path = os.path.join(cur_dir, 'dict/movie_time.txt')
        # 加载特征词
        self.director_words = [i.strip() for i in open(self.director_path, encoding='utf-8') if i.strip()]
        self.actor_words = [i.strip() for i in open(self.actor_path, encoding='utf-8') if i.strip()]
        self.movie_words = [i.strip() for i in open(self.movie_path, encoding='utf-8') if i.strip()]
        self.type_words = [i.strip() for i in open(self.type_path, encoding='utf-8') if i.strip()]
        self.date_words = [i.strip() for i in open(self.date_path, encoding='utf-8') if i.strip()]
        self.star_words = [i.strip() for i in open(self.star_path, encoding='utf-8') if i.strip()]
        self.time_words = [i.strip() for i in open(self.time_path, encoding='utf-8') if i.strip()]
        self.region_words = set(self.director_words + self.actor_words + self.movie_words + self.type_words +
                                self.date_words + self.star_words + self.time_words)
        self.less_words = [i.strip() for i in open(self.less_path, encoding='utf-8') if i.strip()]
        self.more_words = [i.strip() for i in open(self.more_path, encoding='utf-8') if i.strip()]
        self.equal_words = [i.strip() for i in open(self.equal_path, encoding='utf-8') if i.strip()]
        # 构造actree 还是想拿HANLP代替但个别格式化不会（
        self.region_tree = self.build_actree(list(self.region_words))
        # 构建词典
        self.wdtype_dict = self.build_word_type_dict()
        # 电影查信息可能的问句疑问词
        self.star = ['评分', '是几', '分数', '有多少', '打分', '打分是多少', '评分是多少', '分数是多少',
                     '有多少分', '有多少分数', '有多少评分', '有多少星', '有多少星级']
        self.time = ['时长', '电影时长', '有多长时间', '时长是多长', '时间是', '时间为', '有多长', '多长']
        self.show_start = ['上映时间', '是什么时候', '什么时候上映', '首映', '在影院上线', '首播',
                           '什么时候可以在影院看到', '在影院放映', '在影院上映', '日期', '上线', '上映', '上映日期']
        self.director_gwds = ['导演是', '编导是', '谁是导演', '谁是编导', '导演']
        self.movie_type = ['类型是', '哪类', '类别是', '的类别', '的类型', '电影类别', '电影类型', '类型的', '类别的',
                           '类型']
        self.actor_gwds = ['主演是', '演员', '谁演的', '演出', '演为', '主演有', '主演都有']
        # 通过信息查电影可能的疑问词
        self.movie_toMovie = ['电影有哪些', '的电影', '电影有']
        self.type_movie = ['类的电影', '类电影有', '类型的电影', '类别的电影']
        self.actor_movie = ['演了', '主演了', '饰演了', '扮演了', '出镜', '主演的', '参演']
        self.director_movie = ['导演了', '拍摄了', '拍了', '主导', '导演的']
        self.associate_movie = ['合作的电影', '合作过的', '一起演了', '一起演过', '合作了', '共同主演', '共同演了',
                                '一起演的', '一块演的', '共同出演', '一齐出演', '一起出演', '合拍的']
        self.associate_person = ['合作的演员', '搭档', '与', '合作了', '共同主演', '共同演了', '一起演的电影',
                                 '合作过的',
                                 '一起演的', '一块演的', '一块演了', '共同出演', '一起演过', '一起出演', '合拍的']
        # self.date_range = ['之后的', '之前的', '之后上映', '之前上映', '以后的', '以前的', '以后上映', '以前上映']
        # self.star_range = ['分以上', '分以下', '大', '小', '分']
        # self.time_range = ['时长', '分钟以', '大于', '小于', '之内', '之外', '以上', '以下', '时间']
        self.date = ['年', '月', '日', '年份', '日期', '时间', '年代', '上映']
        self.star = ['分', '星', '评分', '分数', '星级']
        self.time = ['时长', '时间', '分钟', '小时']
        print('初始化完成')
        return

    '''分类主函数'''

    def classify(self, question):
        data = {}
        movie_dict = self.check_movies(question)
        if not movie_dict:
            return {}
        data['args'] = movie_dict
        # 收集问句当中所涉及到的实体类型
        types = []
        for type_ in movie_dict.values():
            types += type_
        question_types = []

        # 评分有多少
        if self.check_words(self.star, question) and ('movie' in types):
            question_type = 'movie_star'
            question_types.append(question_type)

        # 时长有多长
        if self.check_words(self.time, question) and ('movie' in types):
            question_type = 'movie_time'
            question_types.append(question_type)

        # 首映时间
        if self.check_words(self.show_start, question) and ('movie' in types):
            question_type = 'movie_show_start'
            question_types.append(question_type)

        # 导演是谁
        if self.check_words(self.director_gwds, question) and ('movie' in types):
            question_type = 'movie_director'
            question_types.append(question_type)

        # 主演是谁
        if self.check_words(self.actor_gwds, question) and ('movie' in types):
            question_type = 'movie_actor'
            question_types.append(question_type)

        # 电影类型
        if self.check_words(self.movie_type, question) and ('movie' in types):
            question_type = 'movie_type'
            question_types.append(question_type)

        # 根据其它信息查电影
        if self.check_words(self.movie_toMovie, question) and ('movie' in types):
            question_type = 'movie_toMovie'
            question_types.append(question_type)

        if self.check_words(self.type_movie, question) and ('type' in types):
            question_type = 'type_movie'
            question_types.append(question_type)

        if self.check_words(self.actor_movie, question) and ('actor' in types):
            question_type = 'actor_movie'
            question_types.append(question_type)

        if self.check_words(self.director_movie, question) and ('director' in types):
            question_type = 'director_movie'
            question_types.append(question_type)

        if self.check_words(self.date, question) and ('date' in types):
            question_type = 'date_movie'
            List = []
            for i in movie_dict:
                if i.isdigit:
                    List.append(i)
            if len(List) == 1:
                less = self.check_words(self.less_words, question)
                more = self.check_words(self.more_words, question)
                equal = self.check_words(self.equal_words, question)
                if less:
                    question_type = 'less_date_movie'
                elif more:
                    question_type = 'more_date_movie'
                elif equal:
                    question_type = 'date_movie'
            else:
                question_type = 'range_date_movie'
            question_types.append(question_type)

        if self.check_words(self.time, question) and ('time' in types):
            question_type = 'time_movie'
            List = []
            for i in movie_dict:
                if i.isdigit():
                    List.append(i)
            if len(List) == 1:
                less = self.check_words(self.less_words, question)
                more = self.check_words(self.more_words, question)
                equal = self.check_words(self.equal_words, question)
                if less:
                    question_type = 'less_time_movie'
                elif more:
                    question_type = 'more_time_movie'
                elif equal:
                    question_type = 'time_movie'
            else:
                question_type = 'range_time_movie'
            question_types.append(question_type)

        if self.check_words(self.star, question) and ('star' in types):
            question_type = 'star_movie'
            List = []
            for i in movie_dict:
                try:
                    I = float(i)
                    List.append(I)
                except:
                    pass
            if len(List) == 1:
                less = self.check_words(self.less_words, question)
                more = self.check_words(self.more_words, question)
                equal = self.check_words(self.equal_words, question)
                if less:
                    question_type = 'less_star_movie'
                elif more:
                    question_type = 'more_star_movie'
                elif equal:
                    question_type = 'star_movie'
            else:
                question_type = 'range_star_movie'
            question_types.append(question_type)

        # 联合查询电影
        if self.check_words(self.associate_movie, question) and ('actor' in types):
            question_type = 'associate_movie'
            question_types.append(question_type)

        if self.check_words(self.associate_person, question) and ('actor' in types):
            question_type = 'associate_person'
            question_types.append(question_type)

        data['question_types'] = question_types
        return data

    '''构造词对应的类型'''

    def build_word_type_dict(self):
        wd_dict = dict()
        for wd in self.region_words:
            wd_dict[wd] = []
            if wd in self.actor_words:
                wd_dict[wd].append('actor')
            if wd in self.director_words:
                wd_dict[wd].append('director')
            if wd in self.movie_words:
                wd_dict[wd].append('movie')
            if wd in self.type_words:
                wd_dict[wd].append('type')
            if wd in self.date_words:
                wd_dict[wd].append('date')
            if wd in self.star_words:
                wd_dict[wd].append('star')
            if wd in self.time_words:
                wd_dict[wd].append('time')
        return wd_dict

    '''构造actree，加速过滤'''

    def build_actree(self, wordlist):
        actree = ahocorasick.Automaton()
        for index, word in enumerate(wordlist):
            actree.add_word(word, (index, word))
        actree.make_automaton()
        return actree

    '''问句过滤'''

    def check_movies(self, question):
        region_wds = []
        for i in self.region_tree.iter(question):
            wd = i[1][1]
            region_wds.append(wd)
        stop_wds = []
        for wd1 in region_wds:
            for wd2 in region_wds:
                if wd1 in wd2 and wd1 != wd2:
                    stop_wds.append(wd1)
        final_wds = [i for i in region_wds if i not in stop_wds]
        final_dict = {i: self.wdtype_dict.get(i) for i in final_wds}
        return final_dict

    '''基于特征词进行分类'''

    def check_words(self, wds, sent):
        for wd in wds:
            if wd in sent:
                return True
        return False


if __name__ == '__main__':
    handler = QuestionClassifier()
    while 1:
        question = input('input an question:')
        data = handler.classify(question)
        print(data)
